// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/aggregates/approx_percentile_aggregate.h>
#include <pollux/common/base/io_utils.h>
#include <pollux/common/base/macros.h>
#include <pollux/common/base/random_util.h>
#include <pollux/common/memory/hash_string_allocator.h>
#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/lib/kll_sketch.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        template<typename T, typename Allocator>
        struct KllSketchTypeTraits {
            using KllSketchType = functions::kll::KllSketch<T, Allocator, std::less<T> >;
        };

        template<>
        struct KllSketchTypeTraits<float, StlAllocator<float> > {
            using KllSketchType = functions::kll::KllSketch<
                float,
                StlAllocator<float>,
                util::floating_point::NaNAwareLessThan<float> >;
        };

        template<>
        struct KllSketchTypeTraits<double, StlAllocator<double> > {
            using KllSketchType = functions::kll::KllSketch<
                double,
                StlAllocator<double>,
                util::floating_point::NaNAwareLessThan<double> >;
        };

        template<>
        struct KllSketchTypeTraits<float, std::allocator<float> > {
            using KllSketchType = functions::kll::KllSketch<
                float,
                std::allocator<float>,
                util::floating_point::NaNAwareLessThan<float> >;
        };

        template<>
        struct KllSketchTypeTraits<double, std::allocator<double> > {
            using KllSketchType = functions::kll::KllSketch<
                double,
                std::allocator<double>,
                util::floating_point::NaNAwareLessThan<double> >;
        };

        template<typename T, typename Allocator = StlAllocator<T> >
        using KllSketch = typename KllSketchTypeTraits<T, Allocator>::KllSketchType;
        template<typename T>
        using KllView = functions::kll::detail::View<T>;

        unsigned getRandomSeed(std::optional<uint32_t> fixedRandomSeed) {
            return fixedRandomSeed.has_value() ? *fixedRandomSeed : random::getSeed();
        }

        // Accumulator to buffer large count values in addition to the KLL
        // sketch itself.
        template<typename T>
        struct KllSketchAccumulator {
            explicit KllSketchAccumulator(
                HashStringAllocator *allocator,
                std::optional<uint32_t> fixedRandomSeed)
                : sketch_(
                      functions::kll::kDefaultK,
                      StlAllocator<T>(allocator),
                      getRandomSeed(fixedRandomSeed)),
                  largeCountValues_(StlAllocator<std::pair<T, int64_t> >(allocator)) {
            }

            void setAccuracy(double value) {
                sketch_.setK(functions::kll::kFromEpsilon(value));
            }

            void append(T value) {
                sketch_.insert(value);
            }

            void append(
                T value,
                int64_t count,
                HashStringAllocator *allocator,
                std::optional<uint32_t> fixedRandomSeed) {
                constexpr size_t kMaxBufferSize = 4096;
                constexpr int64_t kMinCountToBuffer = 512;
                if (count < kMinCountToBuffer) {
                    for (int i = 0; i < count; ++i) {
                        sketch_.insert(value);
                    }
                } else {
                    largeCountValues_.emplace_back(value, count);
                    if (largeCountValues_.size() >= kMaxBufferSize) {
                        flush(allocator, fixedRandomSeed);
                    }
                }
            }

            void append(const KllView<T> &view) {
                sketch_.mergeViews(melon::Range(&view, 1));
            }

            void append(const std::vector<KllView<T> > &views) {
                sketch_.mergeViews(views);
            }

            // Creates a copy of the KllSketch, merges the largeCountValues_ into it,
            // compacts it, and returns it.
            // Makes a copy so that uses std::allocator so that this is safe to call
            // during spilling which may run in parallel.  HashStringAllocator is not
            // thread safe, so merging into/compacting the original KllSketch which
            // depends on it can lead to concurrency bugs.
            KllSketch<T, std::allocator<T> > compact(
                std::optional<uint32_t> fixedRandomSeed) const {
                KllSketch<T, std::allocator<T> > newSketch =
                        KllSketch<T, std::allocator<T> >::fromView(
                            sketch_.toView(),
                            std::allocator<T>(),
                            getRandomSeed(fixedRandomSeed));

                mergeLargeCountValuesIntoSketch(
                    std::allocator<T>(), newSketch, fixedRandomSeed);

                newSketch.compact();

                return newSketch;
            }

            const KllSketch<T> &getSketch() const {
                return sketch_;
            }

            // This must be called before the KllSketch can be used for estimateQuantile()
            // or estimateQuantiles().
            void flush(
                HashStringAllocator *allocator,
                std::optional<uint32_t> fixedRandomSeed) {
                mergeLargeCountValuesIntoSketch(
                    StlAllocator<T>(allocator), sketch_, fixedRandomSeed);
                largeCountValues_.clear();

                sketch_.finish();
            }

        private:
            template<typename Allocator, typename Compare>
            void mergeLargeCountValuesIntoSketch(
                const Allocator &allocator,
                functions::kll::KllSketch<T, Allocator, Compare> &sketch,
                std::optional<uint32_t> fixedRandomSeed) const {
                if (!largeCountValues_.empty()) {
                    std::vector<functions::kll::KllSketch<T, Allocator, Compare> > sketches;
                    sketches.reserve(largeCountValues_.size());
                    for (auto [x, n]: largeCountValues_) {
                        sketches.push_back(
                            functions::kll::KllSketch<T, Allocator, Compare>::fromRepeatedValue(
                                x, n, sketch_.k(), allocator, getRandomSeed(fixedRandomSeed)));
                    }
                    sketch.merge(melon::Range(sketches.begin(), sketches.end()));
                }
            }

            KllSketch<T> sketch_;
            std::vector<std::pair<T, int64_t>, StlAllocator<std::pair<T, int64_t> > >
            largeCountValues_;
        };

        void checkWeight(int64_t weight) {
            constexpr int64_t kMaxWeight = (1ll << 60) - 1;
            POLLUX_USER_CHECK(
                1 <= weight && weight <= kMaxWeight,
                "{}: weight must be in range [1, {}], got {}",
                kApproxPercentile,
                kMaxWeight,
                weight);
        }

        template<typename T>
        class ApproxPercentileAggregate : public exec::Aggregate {
        public:
            ApproxPercentileAggregate(
                bool hasWeight,
                bool hasAccuracy,
                const TypePtr &resultType,
                std::optional<uint32_t> fixedRandomSeed)
                : exec::Aggregate(resultType),
                  hasWeight_{hasWeight},
                  hasAccuracy_(hasAccuracy),
                  fixedRandomSeed_(fixedRandomSeed) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(KllSketchAccumulator<T>);
            }

            bool isFixedSize() const override {
                return false;
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                for (auto i = 0; i < numGroups; ++i) {
                    value<KllSketchAccumulator<T> >(groups[i])->flush(
                        allocator_, fixedRandomSeed_);
                }

                POLLUX_CHECK(result);
                // When all inputs are nulls or masked out, percentiles_ can be
                // uninitialized. The result should be nulls in this case.
                if (!percentiles_.has_value()) {
                    *result = BaseVector::create_null_constant(
                        (*result)->type(), numGroups, (*result)->pool());
                    return;
                }

                if (percentiles_ && percentiles_->is_array) {
                    melon::Range percentiles(
                        percentiles_->values.begin(), percentiles_->values.end());
                    auto arrayResult = (*result)->as_unchecked<ArrayVector>();
                    vector_size_t elementsCount = 0;
                    for (auto i = 0; i < numGroups; ++i) {
                        char *group = groups[i];
                        auto accumulator = value<KllSketchAccumulator<T> >(group);
                        if (accumulator->getSketch().totalCount() > 0) {
                            elementsCount += percentiles.size();
                        }
                    }
                    arrayResult->elements()->resize(elementsCount);
                    elementsCount = 0;
                    auto rawValues =
                            arrayResult->elements()->as_flat_vector<T>()->mutableRawValues();
                    extract(
                        groups,
                        numGroups,
                        arrayResult,
                        [&](const KllSketch<T> &digest,
                            ArrayVector *result,
                            vector_size_t index) {
                            digest.estimateQuantiles(percentiles, rawValues + elementsCount);
                            result->setOffsetAndSize(index, elementsCount, percentiles.size());
                            result->set_null(index, false);
                            elementsCount += percentiles.size();
                        });
                } else {
                    extract(
                        groups,
                        numGroups,
                        (*result)->as_flat_vector<T>(),
                        [&](const KllSketch<T> &digest,
                            FlatVector<T> *result,
                            vector_size_t index) {
                            POLLUX_DCHECK_EQ(percentiles_->values.size(), 1);
                            result->set(
                                index, digest.estimateQuantile(percentiles_->values.back()));
                        });
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                std::vector<KllSketch<T, std::allocator<T> > > sketches;
                sketches.reserve(numGroups);
                for (auto i = 0; i < numGroups; ++i) {
                    sketches.push_back(
                        value<KllSketchAccumulator<T> >(groups[i])->compact(fixedRandomSeed_));
                }

                POLLUX_CHECK(result);
                auto rowResult = (*result)->as<RowVector>();
                POLLUX_CHECK(rowResult);
                auto pool = rowResult->pool();

                // percentiles_ can be uninitialized during an intermediate aggregation step
                // when all input intermediate states are nulls. Result should be nulls in
                // this case.
                if (!percentiles_) {
                    rowResult->ensure_writable(SelectivityVector{numGroups});
                    // rowResult->childAt(i) for i = kPercentiles, kPercentilesIsArray, and
                    // kAccuracy are expected to be constant in addIntermediateResults.
                    rowResult->childAt(kPercentiles) =
                            BaseVector::create_null_constant(ARRAY(DOUBLE()), numGroups, pool);
                    rowResult->childAt(kPercentilesIsArray) =
                            BaseVector::create_null_constant(BOOLEAN(), numGroups, pool);
                    rowResult->childAt(kAccuracy) =
                            BaseVector::create_null_constant(DOUBLE(), numGroups, pool);

                    // Set nulls for all rows.
                    auto rawNulls = rowResult->mutable_raw_nulls();
                    bits::fillBits(rawNulls, 0, rowResult->size(), bits::kNull);
                    return;
                }
                auto &values = percentiles_->values;
                auto size = values.size();
                auto elements =
                        BaseVector::create<FlatVector<double> >(DOUBLE(), size, pool);
                std::copy(values.begin(), values.end(), elements->mutableRawValues());
                auto array = std::make_shared<ArrayVector>(
                    pool,
                    ARRAY(DOUBLE()),
                    nullptr,
                    1,
                    AlignedBuffer::allocate<vector_size_t>(1, pool, 0),
                    AlignedBuffer::allocate<vector_size_t>(1, pool, size),
                    std::move(elements));
                rowResult->childAt(kPercentiles) =
                        BaseVector::wrap_in_constant(numGroups, 0, std::move(array));
                rowResult->childAt(kPercentilesIsArray) =
                        std::make_shared<ConstantVector<bool> >(
                            pool,
                            numGroups,
                            false,
                            BOOLEAN(),
                            static_cast<bool &&>(percentiles_->is_array));
                rowResult->childAt(kAccuracy) = std::make_shared<ConstantVector<double> >(
                    pool,
                    numGroups,
                    accuracy_ == kMissingNormalizedValue,
                    DOUBLE(),
                    static_cast<double &&>(accuracy_));
                auto k = rowResult->childAt(kK)->as_flat_vector<int32_t>();
                auto n = rowResult->childAt(kN)->as_flat_vector<int64_t>();
                auto minValue = rowResult->childAt(kMinValue)->as_flat_vector<T>();
                auto maxValue = rowResult->childAt(kMaxValue)->as_flat_vector<T>();
                auto items = rowResult->childAt(kItems)->as<ArrayVector>();
                auto levels = rowResult->childAt(kLevels)->as<ArrayVector>();

                rowResult->resize(numGroups);
                k->resize(numGroups);
                n->resize(numGroups);
                minValue->resize(numGroups);
                maxValue->resize(numGroups);
                items->resize(numGroups);
                levels->resize(numGroups);

                auto itemsElements = items->elements()->as_flat_vector<T>();
                auto levelsElements = levels->elements()->as_flat_vector<int32_t>();
                size_t itemsCount = 0;
                vector_size_t levelsCount = 0;
                for (auto &sketch: sketches) {
                    auto v = sketch.toView();
                    itemsCount += v.items.size();
                    levelsCount += v.levels.size();
                }
                POLLUX_CHECK_LE(itemsCount, std::numeric_limits<vector_size_t>::max());
                itemsElements->reset_nulls();
                itemsElements->resize(itemsCount);
                levelsElements->reset_nulls();
                levelsElements->resize(levelsCount);

                auto rawItems = itemsElements->mutableRawValues();
                auto rawLevels = levelsElements->mutableRawValues();
                itemsCount = 0;
                levelsCount = 0;
                for (int i = 0; i < sketches.size(); ++i) {
                    auto v = sketches[i].toView();
                    if (v.n == 0) {
                        rowResult->set_null(i, true);
                    } else {
                        rowResult->set_null(i, false);
                        k->set(i, v.k);
                        n->set(i, v.n);
                        minValue->set(i, v.minValue);
                        maxValue->set(i, v.maxValue);
                        std::copy(v.items.begin(), v.items.end(), rawItems + itemsCount);
                        items->setOffsetAndSize(i, itemsCount, v.items.size());
                        itemsCount += v.items.size();
                        std::copy(v.levels.begin(), v.levels.end(), rawLevels + levelsCount);
                        levels->setOffsetAndSize(i, levelsCount, v.levels.size());
                        levelsCount += v.levels.size();
                    }
                }
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                decodeArguments(rows, args);

                if (hasWeight_) {
                    rows.applyToSelected([&](auto row) {
                        if (decodedValue_.is_null_at(row) || decodedWeight_.is_null_at(row)) {
                            return;
                        }

                        auto tracker = trackRowSize(groups[row]);
                        auto accumulator = initRawAccumulator(groups[row]);
                        auto value = decodedValue_.value_at<T>(row);
                        auto weight = decodedWeight_.value_at<int64_t>(row);
                        checkWeight(weight);
                        accumulator->append(value, weight, allocator_, fixedRandomSeed_);
                    });
                } else {
                    if (decodedValue_.may_have_nulls()) {
                        rows.applyToSelected([&](auto row) {
                            if (decodedValue_.is_null_at(row)) {
                                return;
                            }

                            auto accumulator = initRawAccumulator(groups[row]);
                            accumulator->append(decodedValue_.value_at<T>(row));
                        });
                    } else {
                        rows.applyToSelected([&](auto row) {
                            auto accumulator = initRawAccumulator(groups[row]);
                            accumulator->append(decodedValue_.value_at<T>(row));
                        });
                    }
                }
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                addIntermediate<false>(groups, rows, args);
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                decodeArguments(rows, args);

                auto tracker = trackRowSize(group);
                auto accumulator = initRawAccumulator(group);

                if (hasWeight_) {
                    rows.applyToSelected([&](auto row) {
                        if (decodedValue_.is_null_at(row) || decodedWeight_.is_null_at(row)) {
                            return;
                        }

                        auto value = decodedValue_.value_at<T>(row);
                        auto weight = decodedWeight_.value_at<int64_t>(row);
                        checkWeight(weight);
                        accumulator->append(value, weight, allocator_, fixedRandomSeed_);
                    });
                } else {
                    if (decodedValue_.may_have_nulls()) {
                        rows.applyToSelected([&](auto row) {
                            if (decodedValue_.is_null_at(row)) {
                                return;
                            }

                            accumulator->append(decodedValue_.value_at<T>(row));
                        });
                    } else {
                        rows.applyToSelected([&](auto row) {
                            accumulator->append(decodedValue_.value_at<T>(row));
                        });
                    }
                }
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                addIntermediate<true>(group, rows, args);
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                exec::Aggregate::setAllNulls(groups, indices);
                for (auto i: indices) {
                    auto group = groups[i];
                    new(group + offset_)
                            KllSketchAccumulator<T>(allocator_, fixedRandomSeed_);
                }
            }

            void destroyInternal(melon::Range<char **> groups) override {
                for (auto group: groups) {
                    if (isInitialized(group)) {
                        value<KllSketchAccumulator<T> >(group)->~KllSketchAccumulator<T>();
                    }
                }
            }

        private:
            template<typename VectorType, typename ExtractFunc>
            void extract(
                char **groups,
                int32_t numGroups,
                VectorType *result,
                ExtractFunc extractFunction) {
                POLLUX_CHECK(result);
                result->resize(numGroups);

                uint64_t *rawNulls = nullptr;
                if (result->may_have_nulls()) {
                    BufferPtr &nulls = result->mutable_nulls(result->size());
                    rawNulls = nulls->asMutable<uint64_t>();
                }

                for (auto i = 0; i < numGroups; ++i) {
                    char *group = groups[i];
                    auto accumulator = value<KllSketchAccumulator<T> >(group);
                    if (accumulator->getSketch().totalCount() == 0) {
                        result->set_null(i, true);
                    } else {
                        if (rawNulls) {
                            bits::clearBit(rawNulls, i);
                        }
                        extractFunction(accumulator->getSketch(), result, i);
                    }
                }
            }

            void decodeArguments(
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args) {
                size_t argIndex = 0;
                decodedValue_.decode(*args[argIndex++], rows, true);
                if (hasWeight_) {
                    decodedWeight_.decode(*args[argIndex++], rows, true);
                }
                checkSetPercentile(rows, *args[argIndex++]);
                if (hasAccuracy_) {
                    decodedAccuracy_.decode(*args[argIndex++], rows, true);
                    checkSetAccuracy(rows);
                }
                POLLUX_CHECK_EQ(argIndex, args.size());
            }

            void checkSetPercentile(
                const SelectivityVector &rows,
                const BaseVector &vec) {
                DecodedVector decoded(vec, rows);

                auto *base = decoded.base();
                auto baseFirstRow = decoded.index(rows.begin());
                if (!decoded.isConstantMapping()) {
                    rows.applyToSelected([&](vector_size_t row) {
                        POLLUX_USER_CHECK(!decoded.is_null_at(row), "Percentile cannot be null");
                        auto baseRow = decoded.index(row);
                        POLLUX_USER_CHECK(
                            base->equal_value_at(base, baseRow, baseFirstRow),
                            "Percentile argument must be constant for all input rows: {} vs. {}",
                            base->toString(baseRow),
                            base->toString(baseFirstRow));
                    });
                }

                bool is_array;
                vector_size_t offset;
                vector_size_t len;
                if (base->type_kind() == TypeKind::DOUBLE) {
                    is_array = false;
                    offset = rows.begin();
                    len = 1;
                } else if (base->type_kind() == TypeKind::ARRAY) {
                    is_array = true;
                    auto arrays = base->as_unchecked<ArrayVector>();
                    decoded.decode(*arrays->elements());
                    offset = arrays->offsetAt(baseFirstRow);
                    len = arrays->sizeAt(baseFirstRow);
                } else {
                    POLLUX_USER_FAIL(
                        "Incorrect type for percentile: {}", base->type()->toString());
                }
                checkSetPercentile(is_array, decoded, offset, len);
            }

            void checkSetPercentile(
                bool is_array,
                const DecodedVector &percentiles,
                vector_size_t offset,
                vector_size_t len) {
                if (!percentiles_) {
                    POLLUX_USER_CHECK_GT(len, 0, "Percentile cannot be empty");
                    percentiles_ = {
                        .values = std::vector<double>(len),
                        .is_array = is_array,
                    };
                    for (vector_size_t i = 0; i < len; ++i) {
                        POLLUX_USER_CHECK(!percentiles.is_null_at(i), "Percentile cannot be null");
                        auto value = percentiles.value_at<double>(offset + i);
                        POLLUX_USER_CHECK_GE(value, 0, "Percentile must be between 0 and 1");
                        POLLUX_USER_CHECK_LE(value, 1, "Percentile must be between 0 and 1");
                        percentiles_->values[i] = value;
                    }
                } else {
                    POLLUX_USER_CHECK_EQ(
                        is_array,
                        percentiles_->is_array,
                        "Percentile argument must be constant for all input rows");
                    POLLUX_USER_CHECK_EQ(
                        len,
                        percentiles_->values.size(),
                        "Percentile argument must be constant for all input rows");
                    for (vector_size_t i = 0; i < len; ++i) {
                        POLLUX_USER_CHECK_EQ(
                            percentiles.value_at<double>(offset + i),
                            percentiles_->values[i],
                            "Percentile argument must be constant for all input rows");
                    }
                }
            }

            void checkSetAccuracy(const SelectivityVector &rows) {
                if (!hasAccuracy_) {
                    return;
                }

                if (decodedAccuracy_.isConstantMapping()) {
                    POLLUX_USER_CHECK(
                        !decodedAccuracy_.is_null_at(0), "Accuracy cannot be null");
                    checkSetAccuracy(decodedAccuracy_.value_at<double>(0));
                } else {
                    rows.applyToSelected([&](auto row) {
                        POLLUX_USER_CHECK(
                            !decodedAccuracy_.is_null_at(row), "Accuracy cannot be null");
                        const auto accuracy = decodedAccuracy_.value_at<double>(row);
                        if (accuracy_ == kMissingNormalizedValue) {
                            checkSetAccuracy(accuracy);
                        }
                        POLLUX_USER_CHECK_EQ(
                            accuracy,
                            accuracy_,
                            "Accuracy argument must be constant for all input rows");
                    });
                }
            }

            void checkSetAccuracy(double accuracy) {
                POLLUX_USER_CHECK(
                    0 < accuracy && accuracy <= 1, "Accuracy must be between 0 and 1");
                if (accuracy_ == kMissingNormalizedValue) {
                    accuracy_ = accuracy;
                } else {
                    POLLUX_USER_CHECK_EQ(
                        accuracy,
                        accuracy_,
                        "Accuracy argument must be constant for all input rows");
                }
            }

            KllSketchAccumulator<T> *initRawAccumulator(char *group) {
                auto accumulator = value<KllSketchAccumulator<T> >(group);
                if (accuracy_ != kMissingNormalizedValue) {
                    accumulator->setAccuracy(accuracy_);
                }
                return accumulator;
            }

            template<bool kSingleGroup>
            void addIntermediate(
                std::conditional_t<kSingleGroup, char *, char **> group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args) {
                if (validateIntermediateInputs_) {
                    addIntermediateImpl<kSingleGroup, true>(group, rows, args);
                } else {
                    addIntermediateImpl<kSingleGroup, false>(group, rows, args);
                }
            }

            struct Percentiles {
                std::vector<double> values;
                bool is_array;
            };

            static constexpr double kMissingNormalizedValue = -1;
            const bool hasWeight_;
            const bool hasAccuracy_;
            const std::optional<uint32_t> fixedRandomSeed_;
            std::optional<Percentiles> percentiles_;
            double accuracy_{kMissingNormalizedValue};
            DecodedVector decodedValue_;
            DecodedVector decodedWeight_;
            DecodedVector decodedAccuracy_;
            DecodedVector decodedDigest_;

        private:
            template<bool kSingleGroup, bool checkIntermediateInputs>
            void addIntermediateImpl(
                std::conditional_t<kSingleGroup, char *, char **> group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args) {
                POLLUX_CHECK_EQ(args.size(), 1);
                DecodedVector decoded(*args[0], rows);
                auto rowVec = decoded.base()->as<RowVector>();
                if constexpr (checkIntermediateInputs) {
                    POLLUX_USER_CHECK(rowVec);
                    for (int i = kPercentiles; i <= kAccuracy; ++i) {
                        POLLUX_USER_CHECK(rowVec->childAt(i)->is_constant_encoding());
                    }
                    for (int i = kK; i <= kMaxValue; ++i) {
                        POLLUX_USER_CHECK(rowVec->childAt(i)->is_flat_encoding());
                    }
                    for (int i = kItems; i <= kLevels; ++i) {
                        POLLUX_USER_CHECK(
                            rowVec->childAt(i)->encoding() == VectorEncoding::Simple::ARRAY);
                    }
                } else {
                    POLLUX_CHECK(rowVec);
                }

                const SelectivityVector *baseRows = &rows;
                SelectivityVector innerRows{rowVec->size(), false};
                if (!decoded.isIdentityMapping()) {
                    if (decoded.isConstantMapping()) {
                        innerRows.setValid(decoded.index(0), true);
                        innerRows.updateBounds();
                    } else {
                        pollux::translateToInnerRows(
                            rows, decoded.indices(), decoded.nulls(&rows), innerRows);
                    }
                    baseRows = &innerRows;
                }

                DecodedVector percentiles(*rowVec->childAt(kPercentiles), *baseRows);
                auto percentileIsArray =
                        rowVec->childAt(kPercentilesIsArray)->as_unchecked<SimpleVector<bool> >();
                auto accuracy =
                        rowVec->childAt(kAccuracy)->as_unchecked<SimpleVector<double> >();
                auto k = rowVec->childAt(kK)->as_unchecked<SimpleVector<int32_t> >();
                auto n = rowVec->childAt(kN)->as_unchecked<SimpleVector<int64_t> >();
                auto minValue = rowVec->childAt(kMinValue)->as_unchecked<SimpleVector<T> >();
                auto maxValue = rowVec->childAt(kMaxValue)->as_unchecked<SimpleVector<T> >();
                auto items = rowVec->childAt(kItems)->as_unchecked<ArrayVector>();
                auto levels = rowVec->childAt(kLevels)->as_unchecked<ArrayVector>();

                auto itemsElements = items->elements()->as_flat_vector<T>();
                auto levelElements = levels->elements()->as_flat_vector<int32_t>();
                if constexpr (checkIntermediateInputs) {
                    POLLUX_USER_CHECK(itemsElements);
                    POLLUX_USER_CHECK(levelElements);
                } else {
                    POLLUX_CHECK(itemsElements);
                    POLLUX_CHECK(levelElements);
                }
                auto rawItems = itemsElements->rawValues();
                auto rawLevels = levelElements->rawValues<uint32_t>();

                KllSketchAccumulator<T> *accumulator = nullptr;
                std::vector<KllView<T> > views;
                if constexpr (kSingleGroup) {
                    views.reserve(rows.end());
                }
                rows.applyToSelected([&](auto row) {
                    if (decoded.is_null_at(row)) {
                        return;
                    }
                    int i = decoded.index(row);
                    if (percentileIsArray->is_null_at(i)) {
                        return;
                    }
                    if (!accumulator) {
                        int indexInBaseVector = percentiles.index(i);
                        auto percentilesBase = percentiles.base()->as_unchecked<ArrayVector>();
                        auto percentileBaseElements =
                                percentilesBase->elements()->as_flat_vector<double>();
                        if constexpr (checkIntermediateInputs) {
                            POLLUX_USER_CHECK(percentileBaseElements);
                            POLLUX_USER_CHECK(!percentilesBase->is_null_at(indexInBaseVector));
                        }

                        bool is_array = percentileIsArray->value_at(i);
                        DecodedVector decodedElements(*percentilesBase->elements());
                        checkSetPercentile(
                            is_array,
                            decodedElements,
                            percentilesBase->offsetAt(indexInBaseVector),
                            percentilesBase->sizeAt(indexInBaseVector));

                        if (!accuracy->is_null_at(i)) {
                            checkSetAccuracy(accuracy->value_at(i));
                        }
                    }
                    if constexpr (kSingleGroup) {
                        if (!accumulator) {
                            accumulator = initRawAccumulator(group);
                        }
                    } else {
                        accumulator = initRawAccumulator(group[row]);
                    }

                    if constexpr (checkIntermediateInputs) {
                        POLLUX_USER_CHECK(
                            !(k->is_null_at(i) || n->is_null_at(i) || minValue->is_null_at(i) ||
                                maxValue->is_null_at(i) || items->is_null_at(i) ||
                                levels->is_null_at(i)));
                    }
                    KllView<T> v{
                        .k = static_cast<uint32_t>(k->value_at(i)),
                        .n = static_cast<size_t>(n->value_at(i)),
                        .minValue = minValue->value_at(i),
                        .maxValue = maxValue->value_at(i),
                        .items =
                        {
                            rawItems + items->offsetAt(i),
                            static_cast<size_t>(items->sizeAt(i))
                        },
                        .levels =
                        {
                            rawLevels + levels->offsetAt(i),
                            static_cast<size_t>(levels->sizeAt(i))
                        },
                    };
                    if constexpr (kSingleGroup) {
                        views.push_back(v);
                    } else {
                        auto tracker = trackRowSize(group[row]);
                        accumulator->append(v);
                    }
                });
                if constexpr (kSingleGroup) {
                    if (!views.empty()) {
                        auto tracker = trackRowSize(group);
                        accumulator->append(views);
                    }
                }
            }
        };

        bool validPercentileType(const Type &type) {
            if (type.kind() == TypeKind::DOUBLE) {
                return true;
            }
            if (type.kind() != TypeKind::ARRAY) {
                return false;
            }
            return type.as<TypeKind::ARRAY>().elementType()->kind() == TypeKind::DOUBLE;
        }

        void addSignatures(
            const std::string &inputType,
            const std::string &percentileType,
            const std::string &returnType,
            std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > &
            signatures) {
            auto intermediateType = fmt::format(
                "row(array(double), boolean, double, integer, bigint, {0}, {0}, array({0}), array(integer))",
                inputType);
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .returnType(returnType)
                .intermediateType(intermediateType)
                .argumentType(inputType)
                .argumentType(percentileType)
                .build());
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .returnType(returnType)
                .intermediateType(intermediateType)
                .argumentType(inputType)
                .argumentType("bigint")
                .argumentType(percentileType)
                .build());
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .returnType(returnType)
                .intermediateType(intermediateType)
                .argumentType(inputType)
                .argumentType(percentileType)
                .argumentType("double")
                .build());
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .returnType(returnType)
                .intermediateType(intermediateType)
                .argumentType(inputType)
                .argumentType("bigint")
                .argumentType(percentileType)
                .argumentType("double")
                .build());
        }
    } // namespace

    void registerApproxPercentileAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;
        for (const auto &inputType:
             {"tinyint", "smallint", "integer", "bigint", "real", "double"}) {
            addSignatures(inputType, "double", inputType, signatures);
            addSignatures(
                inputType,
                "array(double)",
                fmt::format("array({})", inputType),
                signatures);
        }
        auto name = prefix + kApproxPercentile;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig &config) -> std::unique_ptr<exec::Aggregate> {
                auto isRawInput = exec::isRawInput(step);
                auto hasWeight =
                        argTypes.size() >= 2 && argTypes[1]->kind() == TypeKind::BIGINT;
                bool hasAccuracy = argTypes.size() == (hasWeight ? 4 : 3);
                auto fixedRandomSeed =
                        config.debugAggregationApproxPercentileFixedRandomSeed();

                if (isRawInput) {
                    POLLUX_USER_CHECK_EQ(
                        argTypes.size(),
                        2 + hasWeight + hasAccuracy,
                        "Wrong number of arguments passed to {}",
                        name);
                    if (hasWeight) {
                        POLLUX_USER_CHECK_EQ(
                            argTypes[1]->kind(),
                            TypeKind::BIGINT,
                            "The type of the weight argument of {} must be BIGINT",
                            name);
                    }
                    if (hasAccuracy) {
                        POLLUX_USER_CHECK_EQ(
                            argTypes.back()->kind(),
                            TypeKind::DOUBLE,
                            "The type of the accuracy argument of {} must be DOUBLE",
                            name);
                    }
                    POLLUX_USER_CHECK(
                        validPercentileType(*argTypes[argTypes.size() - 1 - hasAccuracy]),
                        "The type of the percentile argument of {} must be DOUBLE or ARRAY(DOUBLE)",
                        name);
                } else {
                    POLLUX_USER_CHECK_EQ(
                        argTypes.size(),
                        1,
                        "The type of partial result for {} must be ROW",
                        name);
                    POLLUX_USER_CHECK_EQ(
                        argTypes[0]->kind(),
                        TypeKind::ROW,
                        "The type of partial result for {} must be ROW",
                        name);
                }

                TypePtr type;
                if (!isRawInput && exec::isPartialOutput(step)) {
                    type = argTypes[0]->as_row().childAt(kMinValue);
                } else if (isRawInput) {
                    type = argTypes[0];
                } else if (resultType->is_array()) {
                    type = resultType->as<TypeKind::ARRAY>().elementType();
                } else {
                    type = resultType;
                }

                switch (type->kind()) {
                    case TypeKind::TINYINT:
                        return std::make_unique<ApproxPercentileAggregate<int8_t> >(
                            hasWeight, hasAccuracy, resultType, fixedRandomSeed);
                    case TypeKind::SMALLINT:
                        return std::make_unique<ApproxPercentileAggregate<int16_t> >(
                            hasWeight, hasAccuracy, resultType, fixedRandomSeed);
                    case TypeKind::INTEGER:
                        return std::make_unique<ApproxPercentileAggregate<int32_t> >(
                            hasWeight, hasAccuracy, resultType, fixedRandomSeed);
                    case TypeKind::BIGINT:
                        return std::make_unique<ApproxPercentileAggregate<int64_t> >(
                            hasWeight, hasAccuracy, resultType, fixedRandomSeed);
                    case TypeKind::REAL:
                        return std::make_unique<ApproxPercentileAggregate<float> >(
                            hasWeight, hasAccuracy, resultType, fixedRandomSeed);
                    case TypeKind::DOUBLE:
                        return std::make_unique<ApproxPercentileAggregate<double> >(
                            hasWeight, hasAccuracy, resultType, fixedRandomSeed);
                    default:
                        POLLUX_USER_FAIL(
                            "Unsupported input type for {} aggregation {}",
                            name,
                            type->toString());
                }
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
